import os
import json

import numpy as np

if __name__ == '__main__':
    os.environ['CUDA_VISIBLE_DEVICES'] = '1'

import tensorflow as tf
from common import train_data_dir
from transformer_decoder import StockPredictModel

gpus = tf.config.list_physical_devices(device_type='GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

batch_size = 4


def main():
    with open('config.json', 'r', encoding='utf') as file:
        config = json.load(file)

    model = StockPredictModel(params=config)
    model.build(input_shape=[
        (batch_size, 1), (batch_size, 12), (batch_size, 2),
        (batch_size, None, 50), (batch_size, None, 33),
        (batch_size, None, 1)
    ])

    model.load_weights(
        'checkpoint-31-best-val-error-rate-13.772384643554688.hdf5')

    data = np.load(os.path.join(train_data_dir, '中交地产.npz'))

    res = model.predict(
        [data['cpy_id'],
         data['st_ws_one_hot'],
         data['lt_et'],
         data['time_one_hot'],
         data['float_feature'],
         data['mask'],
         ]
    )

    print([each.shape for each in res])
    np.save('result-中交地产.np', res[0])

    print('load success')


main()
